{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8d85202-9f3b-4689-bf80-dee6f7b37e42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torchvision import models\n",
    "def convrelu(in_channels, out_channels, kernel, padding):\n",
    "    return nn.Sequential(\n",
    "        nn.Conv2d(in_channels, out_channels, kernel, padding=padding),\n",
    "        nn.ReLU(inplace=True)\n",
    "    )\n",
    "class ResNetUNet(nn.Module):\n",
    "    def __init__(self, n_class):\n",
    "        super().__init__()\n",
    "\n",
    "        self.base_model = models.resnet18(pretrained=True)\n",
    "        self.base_layers = list(self.base_model.children())\n",
    "\n",
    "        self.layer0 = nn.Sequential(*self.base_layers[:3])\n",
    "        self.layer0_1x1= convrelu(64, 64, 1, 0)\n",
    "        self.layer1 = nn.Sequential(*self.base_layers[3:5])\n",
    "        self.layer1_1x1= convrelu(64, 64, 1, 0)\n",
    "        self.layer2 = self.base_layers[5]\n",
    "        self.layer2_1x1= convrelu(128, 128, 1, 0)\n",
    "        self.layer3 = self.base_layers[6]\n",
    "        self.layer3_1x1= convrelu(256, 256, 1, 0)\n",
    "        self.layer4 = self.base_layers[7]\n",
    "        self.layer4_1x1= convrelu(512, 512, 1, 0)\n",
    "\n",
    "        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n",
    "        self.conv_up3 = convrelu(256 + 512, 512, 3, 1)\n",
    "        self.conv_up2 = convrelu(128 + 512, 256, 3, 1)\n",
    "        self.conv_up1 = convrelu(64 + 256, 256, 3, 1)\n",
    "        self.conv_up0 = convrelu(64 + 256, 128, 3, 1)\n",
    "        self.conv_original_size0 = convrelu(3, 64, 3, 1)\n",
    "        self.conv_original_size1 = convrelu(64, 64, 3, 1)\n",
    "        self.conv_original_size2 = convrelu(64 + 128, 64, 3, 1)\n",
    "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
    "\n",
    "    def forward(self, input):\n",
    "        x_original = self.conv_original_size0(input)\n",
    "        x_original = self.conv_original_size1(x_original)\n",
    "\n",
    "        layer0 = self.layer0(input)\n",
    "        layer1 = self.layer1(layer0)\n",
    "        layer2 = self.layer2(layer1)\n",
    "        layer3 = self.layer3(layer2)\n",
    "        layer4 = self.layer4(layer3)\n",
    "\n",
    "        layer4 = self.layer4_1x1(layer4)\n",
    "        x = self.upsample(layer4)\n",
    "        layer3 = self.layer3_1x1(layer3)\n",
    "        x = torch.cat([x, layer3], dim=1)\n",
    "        x = self.conv_up3(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer2_1x1(layer2)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up2(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer1 = self.layer1_1x1(layer1)\n",
    "        x = torch.cat([x, layer1], dim=1)\n",
    "        x = self.conv_up1(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        layer2 = self.layer0_1x1(layer0)\n",
    "        x = torch.cat([x, layer2], dim=1)\n",
    "        x = self.conv_up0(x)\n",
    "\n",
    "        x = self.upsample(x)\n",
    "        x = torch.cat([x, x_original], dim=1)\n",
    "        x = self.conv_original_size2(x)\n",
    "\n",
    "        out = self.conv_last(x)\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac8ef75f-51e5-42a0-83e6-91eb9d51f3ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from PIL import Image\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "#图像预处理\n",
    "transform= transforms.Compose([\n",
    "    transforms.Resize((256, 256)),\n",
    "    transforms.ToTensor()])\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, img_dir, mask_dir, transform=None):\n",
    "        #图像文件夹路径\n",
    "        self.img_dir = img_dir\n",
    "        #掩码文件夹路径\n",
    "        self.mask_dir = mask_dir\n",
    "        #预处理操作\n",
    "        self.transform = transform\n",
    "        #获取图像文件夹内所有.jpg文件\n",
    "        self.image_filenames = [f for f in os.listdir(img_dir) if f.endswith('.jpg')]\n",
    "\n",
    "    def __len__(self):\n",
    "        #返回数据集样本总数\n",
    "        return len(self.image_filenames)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        #获取文件名\n",
    "        img_name = self.image_filenames[idx]\n",
    "        #将图像和掩码的文件名与对应的目录路径组合，形成完整的文件路径\n",
    "        img_path = os.path.join(self.img_dir, img_name)\n",
    "        mask_path = os.path.join(self.mask_dir, img_name)\n",
    "\n",
    "        #加载图像和掩码，转换图像格式（RGB和灰度格式）\n",
    "        image = Image.open(img_path).convert(\"RGB\")\n",
    "        mask = Image.open(mask_path).convert(\"L\")\n",
    "\n",
    "        #应用预处理\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "            mask = self.transform(mask)\n",
    "            #分割掩码二值化，非零值转换为1.0，零值保持0.0；统一标签格式。适配损失函数\n",
    "            mask = (mask > 0).float()\n",
    "        return image, mask\n",
    "#创建数据集实例\n",
    "train_dataset = MyDataset(img_dir = r\"C:\\Users\\29036\\Desktop\\Car\\imgs\\train\", mask_dir = r\"C:\\Users\\29036\\Desktop\\Car\\masks\\train\", transform = transform)\n",
    "val_dataset = MyDataset(img_dir = r\"C:\\Users\\29036\\Desktop\\Car\\imgs\\val\", mask_dir = r\"C:\\Users\\29036\\Desktop\\Car\\masks\\val\", transform = transform)\n",
    "#创建数据加载器（shuffle=True 训练打乱 不打乱数据）\n",
    "train_loader = DataLoader(train_dataset, batch_size = 8, shuffle = True)\n",
    "val_loader = DataLoader(val_dataset, batch_size = 8, shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "880481d1-9587-4411-bec6-5374c37bfd31",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "# 二分类任务\n",
    "model = ResNetUNet(n_class=1).to(device)\n",
    "# 二元交叉煽损失函数，适用于二分类任务\n",
    "criterion = nn.BCEWithLogitsLoss()  # 损失函数\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)  # 优化器\n",
    "num_epochs = 10  # 训练轮次"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaac505f-6866-4dbd-991e-0767fc54f683",
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(num_epochs):\n",
    "    print('test:', epoch)\n",
    "    model.train()\n",
    "    # 统计损失值\n",
    "    running_loss = 0.0\n",
    "    # 遍历各个批次（图像及其掩码）\n",
    "    for index,(inputs, masks) in enumerate(train_loader):\n",
    "        inputs, masks = inputs.to(device), masks.to(device)\n",
    "        # 清除之前的梯度信息\n",
    "        optimizer.zero_grad()\n",
    "        # 输入数据的模型，得到预测输出\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, masks)\n",
    "        # 反向传播，计算梯度\n",
    "        loss.backward()\n",
    "        # 优化器更新模型参数\n",
    "        optimizer.step()\n",
    "        # 累积损失值\n",
    "        running_loss += loss.item()\n",
    "    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(train_loader):.4f}')\n",
    "\n",
    "    model.eval()\n",
    "    val_loss = 0.0\n",
    "    with torch.no_grad():\n",
    "        for inputs, masks in val_loader:\n",
    "            inputs, masks = inputs.to(device), masks.to(device)\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, masks)\n",
    "            val_loss += loss.item()\n",
    "\n",
    "        print(f\"Validation Loss : {val_loss/len(val_loader):.4f}\")\n",
    "torch.save(model.state_dict(),\"Last_model.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa1583f1-58aa-4cf8-90e4-33823ed5aa9c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5460ce42-76e9-4ee8-8b70-2c03565aaa17",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71b7f8d8-a294-4e3a-8cc9-993e58f100fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b31bb5f-ce0c-405f-b169-61698cb4309a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "AlexNet",
   "language": "python",
   "name": "alexnet"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
